#!/usr/bin/env python3
"""
Subscribe to /unilidar/imu and publish a nav_msgs/Path and a visualization_msgs/Marker
so rviz2 can display the estimated IMU trajectory and current orientation.

This node performs a very simple integration of linear_acceleration (in body frame)
using the IMU orientation to transform to world frame, subtracts a gravity estimate,
and integrates to velocity and position. It's meant as a visualization helper only
(not a production-grade state estimator).

Usage:
  ros2 run <pkg> imu_to_path.py  (or python3 demo/imu_to_path.py)

Topics:
  Subscribes: /unilidar/imu (sensor_msgs/Imu)
  Publishes: /imu_path (nav_msgs/Path), /imu_marker (visualization_msgs/Marker)

"""
import argparse
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Imu
from nav_msgs.msg import Path
from geometry_msgs.msg import PoseStamped, Pose, Point, Quaternion
from visualization_msgs.msg import Marker
import numpy as np


def quat2R(qw, qx, qy, qz):
    # same as in demo scripts
    R = np.array([
        [1 - 2 * (qy ** 2 + qz ** 2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
        [2 * (qx * qy + qw * qz), 1 - 2 * (qx ** 2 + qz ** 2), 2 * (qy * qz - qw * qx)],
        [2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx ** 2 + qy ** 2)]
    ], dtype=np.float64)
    return R


class ImuToPath(Node):
    def __init__(self, use_walltime: bool = False):
        super().__init__('imu_to_path')
        self.sub = self.create_subscription(Imu, '/unilidar/imu', self.cb_imu, 10)
        self.pub_path = self.create_publisher(Path, '/imu_path', 10)
        self.pub_marker = self.create_publisher(Marker, '/imu_marker', 10)
        self.path = Path()
        self.path.header.frame_id = 'map'
        self.seq = 0
        self.last_time = None
        self.v = np.zeros(3, dtype=np.float64)
        self.t = np.zeros(3, dtype=np.float64)
        self.gravity = np.array([0.0, 0.0, 9.80665], dtype=np.float64)
        self.bias_z = None
        self.use_walltime = use_walltime

    def cb_imu(self, msg: Imu):
        # timestamp
        if msg.header.stamp.sec == 0 and msg.header.stamp.nanosec == 0:
            # fallback: use Node time
            now = self.get_clock().now().nanoseconds * 1e-9
        else:
            now = msg.header.stamp.sec + msg.header.stamp.nanosec * 1e-9
        if self.last_time is None:
            dt = 0.01
            self.last_time = now
        else:
            dt = now - self.last_time
            if dt <= 0:
                dt = 0.01
            self.last_time = now

        # orientation
        q = msg.orientation
        R = quat2R(q.w, q.x, q.y, q.z)

        # raw acceleration in body frame
        acc_b = np.array([msg.linear_acceleration.x,
                          msg.linear_acceleration.y,
                          msg.linear_acceleration.z], dtype=np.float64)

        # simple bias estimate for z using first measurement
        if self.bias_z is None:
            # assume sensor at rest initially, estimate bias so that acc_world_z ~= gravity
            acc_w0 = R @ acc_b
            self.bias_z = acc_w0[2] - self.gravity[2]
            self.get_logger().info(f'Estimated initial bias_z={self.bias_z:.4f}')

        # transform to world and remove gravity + bias
        acc_w = R @ acc_b
        acc_w[2] = acc_w[2] - self.gravity[2] - self.bias_z

        # integrate
        self.v += acc_w * dt
        self.t += self.v * dt

        # append to path
        pose = PoseStamped()
        # allow overriding message timestamp with current wall-clock time to avoid
        # RViz message-filter drops when not using simulated /clock
        if self.use_walltime:
            pose.header.stamp = self.get_clock().now().to_msg()
        else:
            pose.header.stamp = msg.header.stamp
        pose.header.frame_id = self.path.header.frame_id
        pose.pose.position = Point(x=float(self.t[0]), y=float(self.t[1]), z=float(self.t[2]))
        pose.pose.orientation = Quaternion(x=q.x, y=q.y, z=q.z, w=q.w)
        self.path.header.stamp = msg.header.stamp
        self.path.poses.append(pose)
        # limit path length
        if len(self.path.poses) > 1000:
            self.path.poses.pop(0)

        self.pub_path.publish(self.path)

        # publish a simple arrow marker for orientation
        marker = Marker()
        marker.header = pose.header
        marker.ns = 'imu'
        marker.id = 0
        marker.type = Marker.ARROW
        marker.action = Marker.ADD
        marker.scale.x = 0.3
        marker.scale.y = 0.05
        marker.scale.z = 0.05
        marker.color.r = 1.0
        marker.color.g = 0.0
        marker.color.b = 0.0
        marker.color.a = 1.0
        marker.pose = pose.pose
        self.pub_marker.publish(marker)


def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('--use-walltime', action='store_true', help='Stamp published Path/Marker with current wall-clock time instead of IMU header stamps')
    parsed = parser.parse_args(argv[1:] if argv is not None else None)
    rclpy.init(args=argv)
    node = ImuToPath(use_walltime=parsed.use_walltime)
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    node.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
